from typing import Any, Callable, Optional, Tuple
import torchvision as tv
from torchvision import transforms
from torch.utils.data import Dataset
import numpy as np
from torch.utils.data import Subset
import timm.data.transforms as timmtrans
from datasets import load_dataset

# apply strong augmentation to dataset for vit training!
transform_dict = {
    "CIFAR10_train": transforms.Compose(
        [
            timmtrans.RandomResizedCropAndInterpolation(
                (224, 224), interpolation="bicubic"
            ),
            transforms.RandomHorizontalFlip(0.5),
            transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            transforms.RandomErasing(
                0.25,
            ),
        ]
    ),
    "CIFAR10_test": transforms.Compose(
        [
            timmtrans.RandomResizedCropAndInterpolation(
                (224, 224), interpolation="bicubic"
            ),
            # transforms.RandomHorizontalFlip(0.5),
            # transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            # transforms.RandomErasing(
            #     0.25,
            # ),
        ]
    ),
    "CIFAR100_train": transforms.Compose(
        [
            # transforms.Resize(224),
            timmtrans.RandomResizedCropAndInterpolation(
                (224, 224), interpolation="bicubic"
            ),
            transforms.RandomHorizontalFlip(0.5),
            transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            transforms.RandomErasing(
                0.25,
            ),
        ]
    ),
    "CIFAR100_test": transforms.Compose(
        [
            transforms.Resize(224),
            # timmtrans.RandomResizedCropAndInterpolation(
            #     (224, 224), interpolation="bicubic"
            # ),
            # transforms.RandomHorizontalFlip(0.5),
            # transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            # transforms.RandomErasing(
            #     0.25,
            # ),
        ]
    ),
    "tiny-imagenet_train": transforms.Compose(
        [
            # transforms.Resize(224),
            timmtrans.RandomResizedCropAndInterpolation(
                (224, 224), interpolation="bicubic"
            ),
            transforms.RandomHorizontalFlip(0.5),
            transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            transforms.RandomErasing(
                0.25,
            ),
        ]
    ),
    "tiny-imagenet_test": transforms.Compose(
        [
            transforms.Resize(224),
            # timmtrans.RandomResizedCropAndInterpolation(
            #     (224, 224), interpolation="bicubic"
            # ),
            # transforms.RandomHorizontalFlip(0.5),
            # transforms.RandAugment(2),
            transforms.ToTensor(),
            transforms.Normalize((0.4850, 0.4560, 0.4060), (0.2290, 0.2240, 0.2250)),
            # transforms.RandomErasing(
            #     0.25,
            # ),
        ]
    ),
}


class tiny_imagenet(tv.datasets.VisionDataset):
    def __init__(
        self,
        root: str,
        train: bool = True,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)

        partition = "train" if train else "valid"
        self.dataset = load_dataset(root, split=partition)
        self.train = train

    def __getitem__(self, index: int) -> Tuple:
        img, target = self.dataset[index]["image"], self.dataset[index]["label"]
        img = img.convert("RGB")
        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, target

    def __len__(self) -> int:
        return len(self.dataset)

    def extra_repr(self) -> str:
        split = "Train" if self.train is True else "Test"
        return f"Split: {split}"


dataset_dict = {
    "CIFAR10": tv.datasets.CIFAR10,
    "CIFAR100": tv.datasets.CIFAR100,
    "tiny-imagenet": tiny_imagenet,
}


def construct_dataset(
    dataset_name: str,
    root: str,
) -> (Dataset, Dataset):
    """Generate datasets according to the given dataset name (support torchvision.Dataset)"""
    # Get the transformations of the specified dataset
    transform_train = transform_dict[dataset_name + "_train"]
    transform_test = transform_dict[dataset_name + "_test"]
    # Load the training and testing dataset
    train_set = dataset_dict[dataset_name](
        root, train=True, download=True, transform=transform_train
    )
    test_set = dataset_dict[dataset_name](
        root, train=False, download=True, transform=transform_test
    )
    return train_set, test_set
